from keras.layers import Dense
from keras.layers import Dropout
from keras.layers import Flatten
from keras.layers.convolutional import MaxPooling2D
from keras.layers.convolutional import Conv2D
from keras.models import Sequential
from keras.models import load_model
from keras.optimizers import SGD
from keras.datasets import mnist
import numpy as np
from keras.utils import np_utils
from matplotlib import pyplot as plt
from keras import backend
backend.set_image_data_format('channels_first')

seed=7
np.random.seed(seed)

(xtrain,ytrain),(xtest,ytest) = mnist.load_data()

xtrain = xtrain.reshape(xtrain.shape[0],1,28,28).astype('float32')      #像素值切换到0-1
xtrain =xtrain /255
xtest = xtest.reshape(xtest.shape[0],1,28,28).astype('float32')
xtest =xtest /255
print(xtrain.shape,ytrain.shape)

ytrain = np_utils.to_categorical(ytrain)                        #独热编码 可有可无
ytest = np_utils.to_categorical(ytest)

def create_model():
    model = Sequential()
    model.add(Conv2D(32,(5,5),input_shape=(1,28,28),activation='relu'))  #卷积层
    model.add(MaxPooling2D(pool_size=(2,2)))                            #最大池化层
    model.add(Dropout(0.2))                                                     #Dropout层
    model.add(Flatten())  #扁平化层，用来连接全连接层和卷积层
    model.add(Dense(units=128,activation='relu'))    #非线性激活函数
    model.add(Dense(units=10,activation='softmax'))   #softmax层

    learningrate = 0.1
    decay =  0.01
    momentum = 0.8

    sgd = SGD(learning_rate=learningrate,decay=decay,momentum=momentum)
    model.compile(loss='categorical_crossentropy',optimizer=sgd,metrics=['mae', 'acc'])
    return  model

model = create_model()

history = model.fit(xtrain,ytrain,epochs=3,batch_size=200)
score = model.evaluate(xtest,ytest,batch_size=200)
print(score[0]*100)

filepath = 'CNN.h5'
model.save(filepath)

plt.plot(history.history['acc'])
plt.show()
plt.plot(history.history['mae'])
plt.show()
